package org.apache.flink.client.cluster;

import java.util.Map;
import java.util.function.Function;

import org.apache.flink.api.common.ConsumerID;
import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.PipelineOptions;
import org.apache.flink.configuration.RestOptions;
import org.apache.flink.configuration.ProviderManagerOptions;
import org.apache.flink.runtime.minicluster.MiniCluster;
import org.apache.flink.runtime.minicluster.MiniClusterConfiguration;
import org.apache.flink.runtime.minicluster.MiniClusterConsumerClient;
import org.apache.flink.runtime.minicluster.RpcServiceSharing;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * mini集群启动器
 * Starts a {@link MiniCluster} . 
 */
public final class MiniClusterStarter {

    private static final Logger LOG = LoggerFactory.getLogger(MiniClusterStarter.class);

    private final Configuration configuration;
    private final Function<? super MiniClusterConfiguration, ? extends MiniCluster>
            miniClusterFactory;

    public static MiniClusterStarter create() {
        return new MiniClusterStarter(new Configuration(), MiniCluster::new);
    }

    public static MiniClusterStarter createWithFactory(
            Configuration configuration,
            Function<? super MiniClusterConfiguration, ? extends MiniCluster> miniClusterFactory) {
        return new MiniClusterStarter(configuration, miniClusterFactory);
    }

    private MiniClusterStarter(
            Configuration configuration,
            Function<? super MiniClusterConfiguration, ? extends MiniCluster> miniClusterFactory) {
        this.configuration = configuration;
        this.miniClusterFactory = miniClusterFactory;
    }

    /** Starts a {@link MiniCluster}. */
    public MiniClusterConsumerClient startCluster(ClassLoader userCodeClassloader) throws Exception {
        MiniClusterConfiguration miniClusterConfig = getMiniClusterConfig();
        MiniCluster miniCluster = miniClusterFactory.apply(miniClusterConfig);
        miniCluster.start();

        return new MiniClusterConsumerClient(
                new ConsumerID(),
                miniCluster,
                userCodeClassloader,
                MiniClusterConsumerClient.ConsumeFinalizationBehavior.SHUTDOWN_CLUSTER);
    }

    private MiniClusterConfiguration getMiniClusterConfig() {
        Configuration configuration = new Configuration(this.configuration);

        if (!configuration.contains(RestOptions.BIND_PORT)) {
            configuration.setString(RestOptions.BIND_PORT, "0");
        }

        int numTaskManagers =
                configuration.getInteger(
                        ConfigConstants.LOCAL_NUMBER_TASK_MANAGER,
                        ConfigConstants.DEFAULT_LOCAL_NUMBER_TASK_MANAGER);

        Map<String, String> overwriteParallelisms =
                configuration.get(PipelineOptions.PARALLELISM_OVERRIDES);

        int numSlotsPerTaskManager =
                configuration
                        .getOptional(ProviderManagerOptions.NUM_TASK_SLOTS)
                        .orElseGet(
                                () ->
                                ProviderManagerOptions.NUM_TASK_SLOTS.defaultValue() );

        return new MiniClusterConfiguration.Builder()
                .setConfiguration(configuration)
                .setNumTaskManagers(numTaskManagers)
                .setRpcServiceSharing(RpcServiceSharing.SHARED)
                .setNumSlotsPerTaskManager(numSlotsPerTaskManager)
                .build();
    }

    public static void shutDownCluster(MiniCluster miniCluster) {
        miniCluster
                .closeAsync()
                .whenComplete(
                        (ignored, throwable) -> {
                            if (throwable != null) {
                                LOG.warn("Shutdown of MiniCluster failed.", throwable);
                            }
                        });
    }
}
